import torch
import torch.nn as nn

from . import part as pt

from .encoder import Encoder
from .decoder import Decoder


class TransformerModel(nn.Module):
    def __init__(
            self,
            vocab_size: int,     # 词表的大小，用来创建 embedding layer
            N: int = 6,          # encoder 和 decoder 中 有多少 single layer
            d_model: int = 512,  # 做词嵌入的时候，一个 token 编码成的向量的长度
            headsNum: int = 8,   # 注意力机制中有多少个 注意力头
            d_ff = 2048,         # encoder 和 decoder 中 feed-forward 的隐藏层神经元个数
            dropout = 0.1        # 所有 的 dropout 层 dropout 的 概率
    ):
        super(TransformerModel, self).__init__()

        self.encoder = Encoder(N, headsNum, d_model, d_ff, dropout)
        self.decoder = Decoder(N, headsNum, d_model, d_ff, dropout)

        # int 输入做编码 （输入和输出共用一个词表和嵌入层）
        self.embed = pt.EmbeddingsLayer(vocab_size, d_model)

        # 位置编码
        self.pos_encoder = pt.PositionEncodingLayer(d_model, dropout=dropout)

        # 对 decoder 的最终输出做 linear 和 softmax
        self.out = pt.GeneratorLayer(d_model, vocab_size)

        # 遍历本模型中所有的参数
        for p in self.parameters():
            # 只对维度大于 1 的参数（比如权重矩阵）做初始化
            if p.dim() > 1:
                # 对参数进行  Xavier 均匀分布初始化
                nn.init.xavier_uniform_(p)


    def encode(
            self,
            src: torch.Tensor,        # (batch_size, src_seq_len) encoder 的输入
            src_mask: torch.Tensor,   # (batch_size, src_seq_len, src_seq_len)
    ) -> torch.Tensor:                # (batch_size, src_seq_len, d_m) 输出输入维度不变

        # 1st, 对 src 词嵌入 + 位置编码
        src = self.pos_encoder(self.embed(src))
        # 2nd, 送入 encoder
        return self.encoder(x=src, mask=src_mask)


    def decode(
            self,
            m: torch.Tensor,           # (batch_size, src_seq_len, d_m) encoder 的 输出

            cross_mask: torch.Tensor,  # (batch_size, tgt_seq_len, src_seq_len)
                                       #   encoder 的输入改的 padding mask

            tgt: torch.Tensor,         # (batch_size, tgt_seq_len, d_m) decoder 的输入

            tgt_mask: torch.Tensor,    # (batch_size, tgt_seq_len, tgt_seq_len)
                                       #   decoder 的输入 padding mask + subsequent mask

    ) -> torch.Tensor:                 # (batch_size, tgt_seq_len, d_m) 和 decoder 输入相同
        # 1st, 对 tgt 词嵌入 + 位置编码
        tgt = self.pos_encoder(self.embed(tgt))
        # 2nd, 将 src & tgt 送入 decoder
        return self.decoder(x=tgt, m=m, cross_mask=cross_mask, tgt_mask=tgt_mask)

    def generate(self, tgt: torch.Tensor) -> torch.Tensor:
        # 经过 linear 和 softmax 的输出
        return self.out(tgt)

    def forward(
            self,
            src: torch.Tensor,         # (batch_size, src_len, d_m) encoder 的输入
            tgt: torch.Tensor,         # (batch_size, tgt_len, d_m) decoder 的输出

            src_mask: torch.Tensor,    # (batch_size, src_len, src_len)
                                       #   encoder 中 self-attention 的 mask

            tgt_mask: torch.Tensor,    # (batch_size, tgt_len, tgt_len)
                                       #   decoder 中 self-attention 的 mask

            cross_mask: torch.Tensor,  # (batch_size, tgt_len, src_len)
                                       #   decoder 中 cross-attention 的 mask

    ) -> torch.Tensor:                 # (batch_size, tgt_len, vocab)

        src = self.encode(src=src, src_mask=src_mask)
        tgt = self.decode(m=src, cross_mask=cross_mask, tgt=tgt, tgt_mask=tgt_mask)
        return self.generate(tgt)


if __name__ == "__main__":
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("Device:", device)

    batch_size = 2
    src_seq_len = 10
    tgt_seq_len = 8
    src_vocab_size = 1000
    tgt_vocab_size = 1000
    d_model = 512

    # 随机构造源句子和目标句子（用词表索引）
    # 在 0 ~ vocab_size 范围内生成 (batch_size, src_seq_len) 大小的 tensor
    src = torch.randint(0, src_vocab_size, (batch_size, src_seq_len)).to(device)
    tgt = torch.randint(0, tgt_vocab_size, (batch_size, tgt_seq_len)).to(device)
    print("Source input shape:", src.shape)
    print("Target input shape:", tgt.shape)

    # 构造掩码 mask
    # TODO: 在 tranformer 中 encoder 和 decoder 交合的那个 attention模块中，
    #  query 来自 decoder，key 和 value 来自 encoder，这种情况下是哪种 mask ?

    # src_mask = pt.subsequentMask(src_seq_len, src_seq_len).to(device)
    # tgt_mask = pt.subsequentMask(tgt_seq_len, tgt_seq_len).to(device)
    src_mask = None
    tgt_mask = None

    # 实例化模型
    model = TransformerModel(src_vocab_size, tgt_vocab_size, d_model=d_model).to(device)

    # 初始化权重，防止全随机导致表现不佳
    def initialize_weights(model):
        for p in model.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)
    initialize_weights(model)

    model.eval()

    # 前向传播
    output = model(src, src_mask, tgt, tgt_mask)

    # TODO: 正确的输出维度是什么 ？
    print("Output shape:", output.shape)
    print("Output sample (first batch, first token, first 10 logits):", output[0, 0, :10])
